import os
from flask import render_template, request, flash, session, redirect, url_for, send_from_directory, \
    make_response, current_app
from . import correct
from .forms import InputForm, LoginForm
from ..models import User


# XSS注入
@correct.route('/XSSc', methods=['GET', 'POST'])
def XSS():
    form = InputForm()
    if request.method == 'POST':
        body = request.form['body']
        return render_template('correct/XSS.html', body=body)
    return render_template('correct/XSS_form.html', form=form)


# CSRF（未完成）
@correct.route('/CSRFc', methods=['GET', 'POST'])
def CSRF():
    form = LoginForm()
    if request.method == 'POST':
        user = User.query.filter_by(username=form.username.data, password=form.password.data).first()
        if user is not None:
            username = request.form['username']
            # password = request.form['password']
            # username_password = username + password
            response = redirect(url_for('correct.transfer'))
            response.set_cookie('username', username)
            return response
        flash('用户名或密码错误')
    return render_template('correct/CSRF.html', form=form)


@correct.route('/transferc', methods=["POST", "GET"])
def transfer():
    # 从cookie中取到用户名
    # username = request.cookies.get('username', None)
    # # 如果没有取到，代表没有登录
    # if not username:
    #     time.sleep(3)
    #     return redirect(url_for('correct.CSRF'))

    if request.method == "POST":
        to_account = request.form.get("to_account")
        money = request.form.get("money")

        # 取出表单中的 csrf_token
        form_csrf_token = request.form.get("csrf_token")
        # 取出 cookie 中的 csrf_token
        cookie_csrf_token = request.cookies.get("csrf_token")
        # 进行对比
        if cookie_csrf_token != form_csrf_token:
            return 'token校验失败，可能是非法操作'

        return render_template('correct/transfer_success.html', to_account=to_account, money=money)
        # return '转账 %s 元到 %s 成功' % (money, to_account)
    response = make_response(render_template('correct/transfer.html'))
    return response


@correct.route('/CSRF_attackc', methods=["POST", "GET"])
def CSRF_attack():
    return render_template('correct/CSRF_attack.html')


# HTTP头注入(未完成)
@correct.route('/HTTP_Headerc', methods=['GET', 'POST'])
def HTTP_Header():
    query = request.args.get('query')
    search_result = User.query.filter_by(username=query)
    if search_result.all():
        return render_template('correct/HTTP_Header.html', res=search_result)
    else:
        flash('未找到相关信息')
        return render_template('correct/HTTP_Header.html')


# 目录遍历
@correct.route('/displayc', methods=['GET', 'POST'])
def _display():
    display_file_list = os.listdir(current_app.config['DISPLAY_PATH'])
    return render_template('correct/display.html', file_list=display_file_list)


@correct.route('/displayc/<filename>', methods=['GET', 'POST'])
def display(filename):
    path = current_app.config['DISPLAY_PATH']
    return send_from_directory(path, filename, as_attachment=False)


# SQL注入
@correct.route('/SQL_injectc', methods=['GET', 'POST'])
def SQL_inject():
    form = LoginForm()
    if request.method == 'POST':
        user = User.query.filter_by(username=form.username.data, password=form.password.data).first()
        if user is not None:
            return render_template('success.html')
        flash('用户名或密码错误')
    return render_template('correct/SQL_inject.html', form=form)


# 浏览可下载文件并下载
@correct.route('/Path_traversalc', methods=['GET', 'POST'])
def Path_traversal():
    download_file_list = os.listdir(current_app.config['DOWNLOAD_PATH'])
    return render_template('correct/Path_traversal_and_download.html', file_list=download_file_list)


@correct.route('/downloadc/<filename>', methods=['GET', 'POST'])
def download(filename):
    # return send_from_directory(current_app.config['DOWNLOAD_PATH'], filename, as_attachment=True)
    path = current_app.config['DOWNLOAD_PATH']
    return send_from_directory(path, filename, as_attachment=True)


@correct.route('/uploadc', methods=['GET', 'POST'])
def upload():
    if request.method == 'POST':
        f = request.files['file']
        if f:
            filename = f.filename
            if not filename.rsplit('.', 1)[1] in current_app.config['ALLOWED_EXTENSIONS']:
                flash('文件格式错误！')
                return redirect(url_for('correct.upload'))
            basepath = os.path.dirname(__file__)  # 当前文件所在路径
            upload_path = os.path.join(basepath, r'..\static\upload\upload_file')
            f.save(os.path.join(upload_path, f.filename))
            flash("上传图片成功!")
            return redirect(url_for('correct.upload'))
        else:
            flash('未选择图片!')
            return redirect(url_for('correct.upload'))
    return render_template('correct/upload.html')


@correct.route('/oplogc', methods=['GET', 'POST'])
def oplog():
    return render_template('index.html')


# 渗透测试安全漏洞之水平越权
@correct.route('/testc', methods=['GET', 'POST'])
def test():
    if request.method == 'POST':
        form = LoginForm()
        if request.method == 'POST':
            user = User.query.filter_by(username=form.username.data, password=form.password.data).first()
            if user is not None:

                username = request.form['username']

                response = redirect(url_for('correct.jurisdiction'))
                response.set_cookie('username', username)

                # 添加校验码
                verification_code = 'mustdefine'
                response.set_cookie('verification_code', verification_code)

                return response
        flash('错误的用户名或密码！')
    return render_template('correct/test.html')


@correct.route('/jurisdictionc', methods=['POST', 'GET'])
def jurisdiction():
    # 从cookie中取到用户名
    username = request.cookies.get('username', None)
    verification_code = request.cookies.get('verification_code', None)

    code = 'mustdefine'
    cookie = 'admin'
    # 如果没有取到，代表没有登录
    if username == cookie and verification_code == code:

        return redirect(url_for('correct._display'))
    return redirect(url_for('correct.test'))

